library(MDEI)
library(splines2)
library(MASS)
library(splines)

options(device="quartz")

set.seed(1)

## Generate data ----
n<-250
treat<-seq(-1,1,length=n)
y.true<-4*pi*treat*sin(2*pi*treat)

errs<-rnorm(n,sd=1)
y<-y.true+errs

deriv.true<-4*pi*sin(2*pi*treat)+8*pi^2*treat*cos(2*pi*treat)

# Generate spline bases, put on a common scale
X.spline<-cbind(MDEI:::bs.me(treat, "treat"))
deriv.X.spline<-cbind(MDEI:::dbs.me(treat))

for(i.scale in 1:ncol(X.spline)){
  sd.curr <- sd(X.spline[,i.scale])
  X.spline[,i.scale] <- X.spline[,i.scale]/sd.curr
  deriv.X.spline[,i.scale] <- deriv.X.spline[,i.scale]/sd.curr
}

# Fit using regression functions in MDEI 
n <- length(y)
p <- ncol(X.spline)

alpha.seq <- seq(max(n * log(p), p), p, length = 100)
s1 <- MDEI:::GCV(y, cbind(1,X.spline), alpha.seq, 0.001)
fits<-cbind(1,X.spline)%*%s1$beta
fits.deriv<-cbind(0,deriv.X.spline)%*%as.vector(s1$beta)
  

## Plot for appendix ----

ylim.true<-range(y)
ylim.true[1]<-ylim.true[1]-.8*diff(ylim.true)


options(device="quartz")
pdf("appendixsplinefig.pdf",h=12,w=10)
par(mfcol=c(3,2),mar=c(3.4,3.2,1.2,.4))

## Top left
plot(treat,y.true,type="n",ylim=ylim.true,xlab="",ylab="")
points(treat,y,pch=19,cex=.2)
lines(treat,y.true,lwd=3)
mtext("Treatment",1,line=2,font=2)
mtext("Outcome",2,line=2,font=2)
arrows(x0=.25,y=-7,x1=.56,y1=-3,length=.1)
text(.25,-7,"True Curve",pos=1, cex=1.35)

## Middle left
plot(treat,y.true,type="n",ylim=ylim.true,xlab="",ylab="")
lines(treat,y.true,lwd=3)
for(i in 2:ncol(X.spline)) lines(treat[-c(1,n)],2.4*X.spline[-c(1,n),i]-25)    
mtext("Treatment",1,line=2,font=2)
mtext("Outcome",2,line=2,font=2)
text(0,-13,"Approximating Basis Functions", cex=1.35)
arrows(x0=-.4,y0=-14,x1=-.6,y1=-17,length=.1)
arrows(x0=.4,y0=-14,x1=.6,y1=-17,length=.1)


## Bottom left
plot(treat,y.true,type="n",ylim=ylim.true,xlab="",ylab="")
lines(treat,y.true,lwd=2,col=gray(.5))
lines(treat,fits,lwd=3)
for(i in which(abs(s1$beta[-1])>0.01)) {
  print(s1$beta[-1][i])
  lines(treat,X.spline[,i]*s1$beta[-1][i]-13,
                                 lwd=3)
}
mtext("Treatment",1,line=2,font=2)
mtext("Outcome",2,line=2,font=2)


points.plot<-c(85,150,239)
points(treat[points.plot],y[points.plot],pch=19,cex=1.5)
arrows(x0=treat[points.plot],y0=y[points.plot],
       x1=treat[points.plot]+.1,y1=y[points.plot]+.1*fits.deriv[points.plot],
       length=.1
       )


text(0,-8+2,"Selected Basis Functions", cex=1.35,pos=3)
arrows(x0=-.32,y0=-8+2,x1=-.45,y1=-9+2,length=.1)
arrows(x0=.32,y0=-8+2,x1=.45,y1=-9+2,length=.1)


arrows(x0=-.7,y0=5,x1=-.5,y1=.5,length=.1)
text(-.7,5,"Fitted Curve",pos=3, cex=1.35)


## Top right
lim.deriv<-range(deriv.true+errs)
lim.deriv[1]<-lim.deriv[1]-1.5*diff(lim.deriv)
plot(treat,deriv.true,type="n",ylim=lim.deriv,xlab="",ylab="")
#points(treat,deriv.true+errs,pch=19,cex=.2)
lines(treat,deriv.true,lwd=3)
mtext("Treatment",1,line=2,font=2)
mtext("Treatment Effect",2,line=2,font=2)
text(-.5,-80,pos=1,"True First Derivative", cex=1.35)
arrows(x0=-.5,y0=-80,x1=-.75,y1=0,length=.1)


## Middle right
plot(treat,deriv.true,type="n",ylim=lim.deriv,xlab="",ylab="")
lines(treat,deriv.true,lwd=3)
for(i in 2:ncol(X.spline)) lines(treat[-c(1,n)],deriv.X.spline[-c(1,n),i]*2.4-200)    
mtext("Treatment",1,line=2,font=2)
mtext("Treatment Effect",2,line=2,font=2)

text(0,-140+20,"Approximating Basis Functions", cex=1.35)
arrows(x0=-.45,y0=-145+15,x1=-.6,y1=-175+15,length=.1)
arrows(x0=.45,y0=-145+15,x1=.6,y1=-175+15,length=.1)

## Bottom right
plot(treat,deriv.true,type="n",ylim=lim.deriv,xlab="",ylab="")
lines(treat,deriv.true,lwd=3,col=gray(.5))
lines(treat,fits.deriv,lwd=3)
for(i in which(abs(s1$beta[-1])>0.01)) {
  #print(s1$coef[-1][i])
  lines(treat,deriv.X.spline[,i]*s1$beta[-1][i]-200,
        lwd=3)
}
mtext("Treatment",1,line=2,font=2)
mtext("Treatment Effect",2,line=2,font=2)

text(0,-120,"Selected Basis Functions", cex=1.35)
arrows(x0=-.32,y0=-130,x1=-.4,y1=-159,length=.1)
arrows(x0=.32,y0=-130,x1=.5,y1=-180,length=.1)


arrows(x0=-.3,y0=-70,x1=-.4,y1=3,length=.1)
text(-.3,-70,"Fitted Derivative Curve",pos=1, cex=1.35)

dev.off()
